import torch.nn as nn
import torch.nn.functional as F
import torch
from torch.autograd import Variable


class LSTMClassifier(nn.Module):

    def __init__(self, embedding_dim, hidden_dim, vocab_size, label_size, use_gpu=False, dropout=0.1):
        super(LSTMClassifier, self).__init__()
        self.hidden_dim = hidden_dim
        self.use_gpu = use_gpu
        self.hidden = None

        self.word_embeddings = nn.Embedding(vocab_size, embedding_dim)
        self.dropout = nn.Dropout(dropout)
        self.lstm = nn.LSTM(embedding_dim, hidden_dim, bidirectional=True)
        # self.fc = nn.Linear(enc_hid_dim * 2, dec_hid_dim)
        self.hidden2label = nn.Linear(hidden_dim*2, label_size)
        # self.hidden = self.init_hidden()

    def init_hidden(self, batch_size, device):
        if self.use_gpu:
            h0 = Variable(torch.zeros(2, batch_size, self.hidden_dim)).to(device)
            c0 = Variable(torch.zeros(2, batch_size, self.hidden_dim)).to(device)
        else:
            h0 = Variable(torch.zeros(2, batch_size, self.hidden_dim))
            c0 = Variable(torch.zeros(2, batch_size, self.hidden_dim))

        self.hidden = (h0, c0)


    def forward(self, src, src_len):
        # assert self.hidden is not None
        # src = [src len, batch size]
        # src_len = [batch size]

        embedded = self.word_embeddings(src)
        embedded = self.dropout(embedded)

        packed_embedded = nn.utils.rnn.pack_padded_sequence(embedded, src_len, enforce_sorted=False)

        packed_outputs, (hidden, _) = self.lstm(packed_embedded)
        outputs, _ = nn.utils.rnn.pad_packed_sequence(packed_outputs)

        # hidden = torch.tanh(self.fc(torch.cat((hidden[-2, :, :], hidden[-1, :, :]), dim=1)))

        # lstm_out, self.hidden = self.lstm(x, self.hidden)
        y = self.hidden2label(torch.cat((hidden[-2, :, :], hidden[-1, :, :]), dim=1))
        genre_scores = F.softmax(y, dim=1)
        return genre_scores

    # calculate acc
    def get_accuracy(self, logits, target, batch_size):
        """ compute accuracy for training round """
        corrects = (torch.max(logits, 1)[1].view(target.size()).data == target.data).sum()
        accuracy = 100.0 * corrects / batch_size
        return accuracy.item()


